import numpy as np
from copy import deepcopy
from collections import defaultdict
import matplotlib.pyplot as plt

def calcCenterPts(dataset, indexs):
    """
    :param dataset: 数据集
    :param indexs: 聚类的索引集合
    :return: centerPts新的中心点
    """
    centerPts = {}
    for k, v in indexs.items():
        centerPts[k] = np.sum(dataset[v, :], axis=0) / len(v)
    return centerPts


def kmeans(dataset, initialPts):
    """
    :param dataset: 数据集
    :param initialPts: 起始点集
    :return: brackets
    """
    flag = True
    tmp = None
    brackets = defaultdict(list)  # 用来装聚类的篮子
    while flag:
        brackets.clear()
        for i in range(len(dataset)):
            min, idx = 65535, 0
            for k, v in initialPts.items():
                distance = np.linalg.norm(v - dataset[i], ord=2)  # 使用欧式距离作为相似度度量
                if min > distance:
                    min = distance
                    idx = k
            brackets[idx].append(i)
        initialPts.clear()
        initialPts = calcCenterPts(dataset, brackets)

        if tmp == brackets:
            flag = False
        tmp = deepcopy(brackets)

    return initialPts, brackets


if __name__ == '__main__':
    colors = ['red', 'blue', 'green', 'orange', 'gray']
    data = np.array([
        [0.4, 0.3],
        [0, 1],
        [2, 2],
        [0, 2],
        [0.13, 0.34],
        [0.3, 0.6],
        [0.4, 0.67],
        [0.4, 0.92],
        [0.61, 0.82],
        [0.63, 0.62],
        [0.58, 0.72],
        [0.41, 0.52],
        [0.74, 0.83],
        [0.81, 1.21],
        [0.71, 0.3],
        [0.78, 0.32],
        [0.81, 0.1],
        [0.2, 0.52],
        [2.1, 0.7],
        [1.3, 1.4],
        [1.5, 1.2],
        [1.45, 1.31],
        [1.61, 1.52],
        [1.47, 1.47],
        [1.4, 1.1],
        [1.3, 1.2],
        [1.7, 1.5],
        [1.2, 1.8],
        [1.5, 1.5],
        [1.5, 1.8],
        [1, 1],
        [0.7, 0.8],
        [0.3, 2.1],
        [1.2, 1.7],
        [2.2, 2.5],
        [0.7, 1.4]
    ])

    initialIdx = np.random.randint(0, len(data), size=3)
    initialPts = {}
    for i in range(len(initialIdx)):
        initialPts[i] = data[initialIdx[i]]
    initPts = deepcopy(initialPts)
    centers, bracket = kmeans(data, initialPts)
    plt.figure(figsize=(14, 6))

    plt.subplot(121)
    plt.scatter(data[:, 0], data[:, 1], c='orange', marker='*', label='samples')
    for k, v in initPts.items():
        plt.scatter(v[0], v[1], c='red', marker='*')
    plt.title('samples and inital points, k=2')
    plt.legend()

    plt.subplot(122)
    for k, v in bracket.items():
        plt.scatter(data[v, 0], data[v, 1], c=colors[k], marker='*', label='class ' + str(k))
    for k, v in centers.items():
        plt.scatter(v[0], v[1], c=colors[k], marker='^', label='CP:class ' + str(k))
    plt.title('k-means')
    plt.legend()

    plt.show()
